import numpy as np
import math
import cv2
from collections import defaultdict

'''关键点list转为dict'''
def landmark_list2dict(landmark):
    '''
    @param: landmark
    将68个关键点按部位保存成dict
    注意下面用的是72个关键点，主要在top_lip，bottom_lip关键点个数的不同（68的有20个，72的有24个），但是不影响正脸校准和关键点的旋转
    'chin': [(290, 56), (257, 114), (233, 173), (216, 233), (206, 292), (203, 359), (210, 418), (223, 474), (259, 519), (312, 547), (371, 550), (433, 543), (491, 520), (544, 486), (592, 441), (635, 391), (668, 340)],
    'left_eyebrow': [(348, 42), (394, 47), (436, 77), (467, 117), (490, 159)],
    'right_eyebrow': [(546, 208), (589, 223), (624, 239), (652, 256), (664, 284)],
    'nose_bridge': [(484, 212), (460, 255), (435, 297), (409, 341)],
    'nose_tip': [(350, 318), (366, 344), (383, 370), (413, 378), (440, 384)],
    'left_eye': [(359, 122), (388, 135), (412, 155), (424, 177), (396, 166), (371, 148)],
    'right_eye': [(537, 255), (573, 265), (598, 287), (613, 308), (583, 301), (557, 280)],
    'top_lip': [(273, 346), (308, 368), (340, 390), (358, 413), (384, 422), (409, 447), (430, 473), (417, 466), (372, 439), (348, 426), (328, 405), (283, 359)],
    'bottom_lip': [(430, 473), (379, 475), (340, 461), (315, 445), (295, 424), (277, 389), (273, 346), (283, 359), (316, 395), (336, 415), (359, 430), (417, 466)]
    '''

    # chin: 1 ~ 17
    # left_eyebrow: 18 ~ 22
    # right_eyebrow: 23 ~ 27
    # nose_bridge: 28 ~ 31
    # nose_tip: 32 ~ 36
    # left_eye: 37 ~ 42
    # right_eye: 43 ~ 48
    # lip: 49 ~ 68
    landmark_list = landmark[0][0]
    landmark_dict = {'part1':[],'left_eye':[],
                     'right_eye':[],'part2':[]}
    for i in range(len(landmark[0][0])):
        if (i < 36):
            landmark_dict['part1'].append(tuple(landmark_list[i]))
        elif (i >= 36 and i < 42):
            landmark_dict['left_eye'].append(tuple(landmark_list[i]))
        elif (i >= 42 and i < 48):
            landmark_dict['right_eye'].append(tuple(landmark_list[i]))
        else:
            landmark_dict['part2'].append(tuple(landmark_list[i]))

    return landmark_dict

'''关键点dict转为list'''
def landmark_dict2list(landmark_dict):
    '''
    :param landmark: dict
    :return:
    '''
    landmark_list = []
    for key in landmark_dict.keys():
        landmark_list.extend(landmark_dict[key])
    return landmark_list

'''
人脸对齐
人脸旋转对齐
人脸对齐思路：
    分别计算左、右眼中心坐标
    计算左右眼中心坐标与水平方向的夹角θ
    计算左右两眼整体中心坐标
    以左右两眼整体中心坐标为基点，将图片array逆时针旋转θ
以下定义了人脸对齐函数
'''
def align_face(image_array, landmarks):
    """ align faces according to eyes position
    :param image_array: numpy array of a single image
    :param landmarks: dict of landmarks for facial parts as keys and tuple of coordinates as values
    :return:
    rotated_img:  numpy array of aligned image
    eye_center: tuple of coordinates for eye center
    angle: degrees of rotation
    """
    # get list landmarks of left and right eye
    left_eye = landmarks['left_eye']
    right_eye = landmarks['right_eye']
    # calculate the mean point of landmarks of left and right eye
    left_eye_center = np.mean(left_eye, axis=0).astype("int")
    right_eye_center = np.mean(right_eye, axis=0).astype("int")
    # compute the angle between the eye centroids
    dy = right_eye_center[1] - left_eye_center[1]
    dx = right_eye_center[0] - left_eye_center[0]
    # compute angle between the line of 2 centeroids and the horizontal line
    angle = math.atan2(dy, dx) * 180. / math.pi
    # calculate the center of 2 eyes
    eye_center = ((left_eye_center[0] + right_eye_center[0]) // 2,
                  (left_eye_center[1] + right_eye_center[1]) // 2)
    # at the eye_center, rotate the image by the angle
    rotate_matrix = cv2.getRotationMatrix2D(eye_center, angle, scale=1)
    rotated_img = cv2.warpAffine(image_array, rotate_matrix, (image_array.shape[1], image_array.shape[0]))
    return rotated_img, eye_center, angle

'''定义旋转图片中landmark的函数，以人脸双眼中心为基点，将每个人脸关键点逆时针旋转θ，该θ角度是人脸对齐的旋转角度。'''
def rotate_landmarks(landmarks, eye_center, angle, row):
    """ rotate landmarks to fit the aligned face
    :param landmarks: dict of landmarks for facial parts as keys and tuple of coordinates as values
    :param eye_center: tuple of coordinates for eye center
    :param angle: degrees of rotation
    :param row: row size of the image
    :return: rotated_landmarks with the same structure with landmarks, but different values
    """
    rotated_landmarks = defaultdict(list)
    for facial_feature in landmarks.keys():
        for landmark in landmarks[facial_feature]:
            rotated_landmark = rotate(origin=eye_center, point=landmark, angle=angle, row=row)
            rotated_landmarks[facial_feature].append(rotated_landmark)
    return rotated_landmarks

'''
人脸关键点旋转
    图片旋转后，图中的landmark坐标也要相应旋转，这样landmark才能匹配旋转后的图片。landmark旋转前的效果如下
'''
'''定义旋转图片中坐标的函数，由于图片和普通坐标系的原点不同，两者坐标点的旋转方式略有出入，图片坐标旋转涉及y坐标在图片坐标系和普通坐标系之间的变换。'''
def rotate(origin, point, angle, row):
    """ rotate coordinates in image coordinate system
    :param origin: tuple of coordinates,the rotation center
    :param point: tuple of coordinates, points to rotate
    :param angle: degrees of rotation
    :param row: row size of the image
    :return: rotated coordinates of point
    """
    x1, y1 = point
    x2, y2 = origin
    y1 = row - y1
    y2 = row - y2
    angle = math.radians(angle)
    x = x2 + math.cos(angle) * (x1 - x2) - math.sin(angle) * (y1 - y2)
    y = y2 + math.sin(angle) * (x1 - x2) + math.cos(angle) * (y1 - y2)
    y = row - y
    return int(x), int(y)
